import torch
from torch.utils.data import DataLoader
from configs.config_setting_tdmaterials import setting_config
from engine import test_one_epoch, train_one_epoch, val_one_epoch
from datasets.dataset_tdmaterials import TDMaterials_dataset, TDMaterial, Split

import os
from tqdm import tqdm

from utils import *

import warnings
warnings.filterwarnings("ignore")


def trainer_tdmaterials(config: setting_config, model, resume_model, checkpoint_dir, logger, writer):
    print('#----------Preparing dataset----------#')
    train_dataset = TDMaterials_dataset(config.data_path, split=Split.TRAIN, transform=config.train_transformer, tdmaterial=TDMaterial.Graphene)
    train_loader = DataLoader(train_dataset,
                                config.batch_size, 
                                shuffle=True,
                                pin_memory=True,
                                num_workers=config.num_workers)
    print("The length of train set is: {}".format(len(train_dataset)))
    val_dataset = TDMaterials_dataset(config.data_path, split=Split.VAL, transform=config.test_transformer, tdmaterial=TDMaterial.Graphene)
    val_loader = DataLoader(val_dataset,
                                batch_size=1,
                                shuffle=False,
                                pin_memory=True, 
                                num_workers=config.num_workers,
                                drop_last=True)
    print("The length of validation set is: {}".format(len(val_dataset)))
    
    max_iterations = config.epochs * len(train_loader)  # max_epoch = max_iterations // len(trainloader) + 1
    logging.info("{} iterations per epoch. {} max iterations ".format(len(train_loader), max_iterations))

    print('#----------Prepareing loss, opt, sch and amp----------#')
    criterion = config.criterion
    optimizer = get_optimizer(config, model)
    scheduler = get_scheduler(config, optimizer)

    print('#----------Set other params----------#')
    min_loss = 999
    start_epoch = 1
    min_epoch = 1

    if os.path.exists(resume_model):
        print('#----------Resume Model and Other params----------#')
        checkpoint = torch.load(resume_model, map_location=torch.device('cpu'))
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        saved_epoch = checkpoint['epoch']
        start_epoch += saved_epoch
        min_loss, min_epoch, loss = checkpoint['min_loss'], checkpoint['min_epoch'], checkpoint['loss']

        log_info = f'resuming model from {resume_model}. resume_epoch: {saved_epoch}, min_loss: {min_loss:.4f}, min_epoch: {min_epoch}, loss: {loss:.4f}'
        logger.info(log_info)

    step = 0
        
    print('#----------Training----------#')
    iterator = tqdm(range(start_epoch, config.epochs + 1), ncols=70)
    for epoch in iterator:

        torch.cuda.empty_cache()

        step = train_one_epoch(
            train_loader,
            model,
            criterion,
            optimizer,
            scheduler,
            epoch,
            step,
            logger,
            config,
            writer
        )

        loss = val_one_epoch(
                val_loader,
                model,
                criterion,
                epoch,
                logger,
                config,
                writer
            )

        if loss < min_loss:
            torch.save(model.state_dict(), os.path.join(checkpoint_dir, 'best.pth'))
            min_loss = loss
            min_epoch = epoch
            
        if epoch!=0 and epoch % config.save_interval == 0:
            torch.save(model.state_dict(), os.path.join(checkpoint_dir, f'epoch{epoch}.pth'))
            
        torch.save(
            {
                'epoch': epoch,
                'min_loss': min_loss,
                'min_epoch': min_epoch,
                'loss': loss,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
            }, os.path.join(checkpoint_dir, 'latest.pth')) 

    if os.path.exists(os.path.join(checkpoint_dir, 'best.pth')):
        print('#----------Testing----------#')
        best_weight = torch.load(config.work_dir + 'checkpoints/best.pth', map_location=torch.device('cpu'))
        model.load_state_dict(best_weight)
        loss = test_one_epoch(
                val_loader,
                model,
                criterion,
                logger,
                config,
            )
        os.rename(
            os.path.join(checkpoint_dir, 'best.pth'),
            os.path.join(checkpoint_dir, f'best-epoch{min_epoch}-loss{min_loss:.4f}.pth')
        )
    